// 导入TensorFlow.js库，机器学习的核心
import * as tf from '@tensorflow/tfjs';
// 导入可视化库，用来画图看数据和训练过程
import * as tfvis from '@tensorflow/tfjs-vis';

// 页面加载完成后执行所有代码
window.onload = async () => {
    // 准备训练数据：身高（厘米）和对应的体重（公斤）
    const heights = [150, 160, 170, 180, 190]
    const weights = [60, 65, 70, 75, 80]
    
    // 画个散点图看看身高体重的分布情况
    tfvis.render.scatterplot(
        { name: '身高体重训练集' }, // 图表标题
        { values: heights.map((h, i) => ({ x: h, y: weights[i] })) }, // 准备数据格式
        { xAxisDomain: [140, 200], yAxisDomain: [50, 90] } // 设置坐标轴范围，看起来舒服点
    )
    
    // 数据归一化处理：把身高缩放到0-1之间，方便模型训练
    // 公式：(身高 - 最小值) / 范围，这里最小值150，范围40（190-150）
    const inputs = tf.tensor(heights).sub(150).div(40)
    // 同样对体重做归一化处理，最小值60，范围20（80-60）
    const labels = tf.tensor(weights).sub(60).div(20)
    
    // 创建一个顺序模型，就是一层一层堆上去的那种
    const model = tf.sequential()
    // 添加一个全连接层，1个输出单元，输入是1维的
    model.add(tf.layers.dense({ units: 1, inputShape: [1] }))

    // 编译模型，用均方误差当损失函数，随机梯度下降当优化器，学习率设为0.1
    model.compile({ loss: tf.losses.meanSquaredError, optimizer: tf.train.sgd(0.1) })
    
    // 开始训练模型啦！
    // 每次喂5个数据（全部数据），训练100轮，顺便可视化训练过程
    await model.fit(inputs, labels, {
        batchSize: 5, // 批次大小
        epochs: 100, // 训练轮数
        callbacks: tfvis.show.fitCallbacks(
            { name: '训练过程' },
            ['loss'], // 只看损失函数的变化就行
        )
    })
    
    // 用训练好的模型预测身高175cm的人体重
    // 注意：这里也要对输入做同样的归一化处理
    const output = model.predict(tf.tensor([175]).sub(150).div(40))
    // 把预测结果反归一化，变回实际体重值
    // 公式：预测值 * 范围 + 最小值
    console.log(output.mul(20).add(60).dataSync()[0])
}